{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10187ae0",
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "num_modes should be even",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[8], line 243\u001b[0m\n\u001b[0;32m    238\u001b[0m             hologram\u001b[38;5;241m.\u001b[39mimg\u001b[38;5;241m.\u001b[39msave(save_path)\n\u001b[0;32m    239\u001b[0m             count \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m--> 243\u001b[0m generate_multimode_advhg_holograms(\u001b[38;5;241m103\u001b[39m, \u001b[38;5;241m5\u001b[39m)\n",
      "Cell \u001b[1;32mIn[8], line 188\u001b[0m, in \u001b[0;36mgenerate_multimode_advhg_holograms\u001b[1;34m(sigma, num_modes, file_folder)\u001b[0m\n\u001b[0;32m    185\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mgenerate_multimode_advhg_holograms\u001b[39m(sigma, num_modes, file_folder\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcurrent folder\u001b[39m\u001b[38;5;124m'\u001b[39m):\n\u001b[0;32m    186\u001b[0m     \u001b[38;5;66;03m# 判断num_modes是否为偶数\u001b[39;00m\n\u001b[0;32m    187\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m num_modes \u001b[38;5;241m%\u001b[39m \u001b[38;5;241m2\u001b[39m \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m--> 188\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnum_modes should be even\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    190\u001b[0m     \u001b[38;5;66;03m# 目标文件夹路径处理\u001b[39;00m\n\u001b[0;32m    191\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m file_folder \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcurrent folder\u001b[39m\u001b[38;5;124m'\u001b[39m:\n",
      "\u001b[1;31mValueError\u001b[0m: num_modes should be even"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from numpy import pi\n",
    "\n",
    "from math import factorial\n",
    "\n",
    "from scipy.special import hermite, laguerre\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "from os import path\n",
    "from PIL import Image\n",
    "\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display, clear_output\n",
    "\n",
    "\n",
    "__all__ = ['SLM', 'HG', 'LG', 'AdvHG', 'CGH']\n",
    "\n",
    "\n",
    "class SLM:\n",
    "    device = 'HoloEye, PLUTO-2-NIR-011'\n",
    "\n",
    "    pixel_size = 8\n",
    "    resolution = (1920, 1080)\n",
    "    \n",
    "    x, y = np.meshgrid(\n",
    "           np.arange(-resolution[0]/2, resolution[0]/2) * pixel_size,\n",
    "          -np.arange(-resolution[1]/2, resolution[1]/2) * pixel_size)\n",
    "    \n",
    "    rho = x**2 + y**2\n",
    "\n",
    "    norm_x = x / (resolution[0] * pixel_size)\n",
    "    norm_y = y / (resolution[1] * pixel_size)\n",
    "\n",
    "\n",
    "\n",
    "class _Mode:\n",
    "    def __init__(self, n, m):\n",
    "        valid = all(isinstance(x, int) and x >= 0 for x in (n, m))\n",
    "        if valid:\n",
    "            self.n = n\n",
    "            self.m = m\n",
    "        else:\n",
    "            raise ValueError('orders must be positive integers')\n",
    "    \n",
    "    @classmethod\n",
    "    def check(self, inputs):\n",
    "        if not isinstance(inputs, (HG, LG, AdvHG, AdvHG_extend)):\n",
    "            raise ValueError('mode must be a instance of HG, LG, AdvHG or AdvHG_extend')\n",
    "\n",
    "\n",
    "\n",
    "class HG(_Mode):\n",
    "    def complex_amplitude(self, w0):\n",
    "        n, m = self.n, self.m\n",
    "        norm = np.sqrt(2**(1-n-m) / (pi * factorial(m) * factorial(n))) / w0\n",
    "        hx, hy= hermite(n)(2**.5 * SLM.x / w0), hermite(m)(2**.5 * SLM.y / w0)\n",
    "        amplitude = norm * hx * hy * np.exp(-SLM.rho/(w0**2))\n",
    "\n",
    "        return amplitude * np.exp(0j)\n",
    "    \n",
    "\n",
    "class AdvHG_extend:\n",
    "    def __init__(self, x1, opx, x2, y1, opy, y2):\n",
    "        import re\n",
    "\n",
    "        if opx not in ('+', '-') or opy not in ('+', '-'):\n",
    "            raise ValueError(\"Operators must be '+' or '-'\")\n",
    "\n",
    "        for val in [x1, x2, y1, y2]:\n",
    "            if not (isinstance(val, int) and val >= 0):\n",
    "                raise ValueError(\"Operands must be non-negative integers\")\n",
    "\n",
    "        self.x1 = x1\n",
    "        self.x2 = x2\n",
    "        self.y1 = y1\n",
    "        self.y2 = y2\n",
    "        self.opx = opx\n",
    "        self.opy = opy\n",
    "\n",
    "        # 运算符映射为数值\n",
    "        op_map = {'+': 1, '-': -1}\n",
    "        self.opx_val = op_map[opx]\n",
    "        self.opy_val = op_map[opy]\n",
    "\n",
    "        mode_x = f\"{x1}{opx}{x2}\"\n",
    "        mode_y = f\"{y1}{opy}{y2}\"\n",
    "        pattern = r'^\\d+[\\+\\-]\\d+$'\n",
    "        if not re.match(pattern, mode_x) or not re.match(pattern, mode_y):\n",
    "            raise ValueError(\"Invalid format: must be like '1+2', '3-1', etc.\")\n",
    "\n",
    "    def complex_amplitude(self, w0):\n",
    "        hgx1y1 = HG(self.x1, self.y1).complex_amplitude(w0)\n",
    "        hgx1y2 = HG(self.x1, self.y2).complex_amplitude(w0)\n",
    "        hgx2y1 = HG(self.x2, self.y1).complex_amplitude(w0)\n",
    "        hgx2y2 = HG(self.x2, self.y2).complex_amplitude(w0)\n",
    "\n",
    "        return (hgx1y1 + self.opy_val * hgx1y2 + self.opx_val * hgx2y1 + self.opx_val * self.opy_val * hgx2y2) / 2      #e.g., <x|0-1>*<y|2+3> = 1/\\sqrt{2}(<x|0>-<x|1>)*1/\\sqrt{2}(<y|2>+<y|3>)\n",
    "\n",
    "class LG(_Mode):\n",
    "    def __init__(self):\n",
    "        raise NotImplementedError('LG mode is not supported yet')\n",
    "\n",
    "\n",
    "\n",
    "class AdvHG:\n",
    "    def __init__(self, pm):\n",
    "        if pm in (-1, 1):\n",
    "            self.pm = pm\n",
    "        else:\n",
    "            raise ValueError('pm must be -1 or 1, 1 stands for + mode, -1 stands for - mode')\n",
    "\n",
    "    def complex_amplitude(self, w0):\n",
    "        hg00 = HG(0, 0).complex_amplitude(w0)\n",
    "        hg01 = HG(0, 1).complex_amplitude(w0)\n",
    "        return (hg00 + self.pm * hg01) / 2**.5\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class CGH:\n",
    "    def __init__(self, sigma, *modes, nx=500, ny=50):\n",
    "\n",
    "        [_Mode.check(mode) for mode in modes]\n",
    "\n",
    "        w0 = 2*sigma\n",
    "\n",
    "        if len(modes) == 1:\n",
    "            ca = modes[0].complex_amplitude(w0)*np.exp(2j*pi*SLM.norm_x*nx)\n",
    "\n",
    "        elif len(modes) == 2:\n",
    "            ca0, ca1 = [m.complex_amplitude(w0) for m in modes]\n",
    "            ca = ca0*np.exp(2j*pi*SLM.norm_y*ny)*np.exp(2j*pi*SLM.norm_x*nx) + ca1*np.exp(-2j*pi*SLM.norm_y*ny)*np.exp(2j*pi*SLM.norm_x*nx)\n",
    "          \n",
    "        elif len(modes) > 2:\n",
    "            # 手动输入参数（角度制）\n",
    "            radius = float(input(\"请输入圆弧的半径：\"))\n",
    "            flareangle = float(input(\"请输入圆弧的张角（单位：度）：\"))\n",
    "            flareangle_rad = np.deg2rad(flareangle)  # 角度 -> 弧度\n",
    "            _nx, _ny = generate_arc_coordinates(radius, flareangle_rad, len(modes))\n",
    "            ca = sum(m.complex_amplitude(w0) *np.exp(2j * pi * SLM.norm_y * y) *np.exp(2j * pi * SLM.norm_x * x) for m, x, y in zip(modes, _nx, _ny))\n",
    "\n",
    "        else:\n",
    "            raise ValueError('invalid modes parameter, possibly due to ' +\n",
    "                             'the number of modes being <= 0')\n",
    "   \n",
    "\n",
    "        \n",
    "        fx2 = interp1d(np.linspace(0, 1, 801),np.load(path.join(os.getcwd(), 'fx2.npy')))\n",
    "        a = np.abs(ca) / np.abs(ca).max()\n",
    "        phi = np.angle(ca)\n",
    "\n",
    "        _temp = fx2(a) * np.sin(phi)\n",
    "        \n",
    "        _temp = ((_temp - _temp.min()) / (_temp.max() - _temp.min())) * 255\n",
    "\n",
    "        self.cgh = _temp.astype(np.uint8)\n",
    "        self.img = Image.fromarray(self.cgh)\n",
    "\n",
    "    def result(self):\n",
    "        return self.cgh\n",
    "    \n",
    "    def show(self):\n",
    "        self.img.show()\n",
    "    \n",
    "    def save(self, file):\n",
    "        file = path.expanduser(file)\n",
    "        # 始终保存（覆盖已有文件）\n",
    "        self.img.save(file)\n",
    "\n",
    "\n",
    "\n",
    "def generate_arc_coordinates(radius, flareangle, num_modes):  #生成等间距圆弧上的坐标点，圆弧位于第一二象限，关于x轴对称\n",
    "    theta = np.linspace(-flareangle / 2, flareangle / 2, num_modes)\n",
    "    _nx = radius * np.cos(theta)\n",
    "    _ny = radius * np.sin(theta)\n",
    "    return _nx, _ny\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#生成一系列全息图，应用于二维多模AdvSPADE成像，例如：当num_modes=10， 则生成100张全息图，分别对应HG(0-1, 0-1), HG(0-1, 0+1), HG(0-1, 2-3), HG(0-1, 2+3)....HG(8+9, 8+9)\n",
    "def generate_multimode_advhg_holograms(sigma, num_modes, file_folder='current folder'):\n",
    "    # 判断num_modes是否为偶数\n",
    "    if num_modes % 2 != 0:\n",
    "        raise ValueError(\"num_modes should be even\")\n",
    "    \n",
    "    # 目标文件夹路径处理\n",
    "    if file_folder == 'current folder':\n",
    "        base_path = os.path.join(os.getcwd(), 'SLM_image')\n",
    "    else:\n",
    "        base_path = os.path.join(os.path.expanduser(file_folder), 'SLM_image')\n",
    "\n",
    "    os.makedirs(base_path, exist_ok=True)\n",
    "\n",
    "    combinations = []\n",
    "\n",
    "    # 构造组合：x=(i±i+1), y=(j±j+1)\n",
    "    for i in range(0, num_modes - 1, 2):\n",
    "        for sign1 in ['-', '+']:\n",
    "            for j in range(0, num_modes - 1, 2):\n",
    "                for sign2 in ['-', '+']:\n",
    "                    combinations.append((i, sign1, i + 1, j, sign2, j + 1))\n",
    "\n",
    "    # 按顺序生成图像并命名\n",
    "    for index, (i, sign1, i2, j, sign2, j2) in enumerate(tqdm(combinations, desc='Generating holograms')):\n",
    "        mode = AdvHG_extend(i, sign1, i2, j, sign2, j2)\n",
    "        hologram = CGH(sigma, mode)\n",
    "\n",
    "        # 文件名格式示例：advhg_000_(0-1,0+1).png\n",
    "        filename = f\"advhg_{index:03d}_({i}{sign1}{i2},{j}{sign2}{j2}).png\"\n",
    "        save_path = os.path.join(base_path, filename)\n",
    "\n",
    "        hologram.save(save_path)\n",
    "\n",
    "\n",
    "#生成一系列全息图，应用于二维多模StdSPADE成像，例如：当num_modes=10， 则生成100张全息图，分别对应HG(0, 1), HG(0, 2), HG(0, 3),....HG(9, 9)\n",
    "def generate_multimode_stdhg_holograms(sigma, num_modes, file_folder='current folder'):\n",
    "    import os\n",
    "    from tqdm import tqdm\n",
    "\n",
    "    if file_folder == 'current folder':\n",
    "        base_path = os.path.join(os.getcwd(), 'SLM_image')\n",
    "    else:\n",
    "        base_path = os.path.join(os.path.expanduser(file_folder), 'SLM_image')\n",
    "\n",
    "    os.makedirs(base_path, exist_ok=True)\n",
    "\n",
    "    count = 1\n",
    "    for i in tqdm(range(num_modes), desc='Generating std HG holograms (rows)'):\n",
    "        for j in range(num_modes):\n",
    "            mode = HG(i, j)\n",
    "            hologram = CGH(sigma, mode)\n",
    "            filename = f\"stdhg_{count:03d}_({i},{j}).png\"\n",
    "            save_path = os.path.join(base_path, filename)\n",
    "            hologram.img.save(save_path)\n",
    "            count += 1\n",
    "\n",
    "\n",
    "\n",
    "generate_multimode_advhg_holograms(103, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f32936d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "w0=206\n",
    "cgh = CGH(103, AdvHG_extend(3, '-', 1, 2, '+', 4))\n",
    "cgh.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dcdd490e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mode1 = HG(2, 0)\n",
    "mode2 = HG(0, 2)\n",
    "mode3 = HG(1, 2)\n",
    "cgh = CGH(103, mode1, mode2, mode3)\n",
    "cgh.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f030d0f8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf",
   "language": "python",
   "name": "tf"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
